import asyncio
import os
from mistralai import Mistral


async def send_message(client, model, messages):
    """
    发送消息并获取流式响应。

    Args:
        client (Mistral): Mistral客户端。
        model (str): 使用的模型名称。
        messages (list): 对话消息列表。

    Returns:
        str: 助手的完整回复。
    """
    # 发送请求并获取流式响应
    response = await client.chat.stream_async(
        model=model,
        messages=messages,
    )
    response_content = ""
    # 处理流式响应
    async for chunk in response:
        if chunk.data.choices[0].delta.content is not None:
            response_content += chunk.data.choices[0].delta.content
            print(chunk.data.choices[0].delta.content, end="")
    return response_content


async def main():
    """
    主函数。
    """
    # 获取环境变量中的API密钥
    api_key = os.environ["MISTRAL_API_KEY"]
    # 设置模型名称
    model = "pixtral-large-latest"

    # 初始化Mistral客户端
    client = Mistral(api_key=api_key)

    # 初始对话历史
    messages = [
        {
            "role": "system",
            "content": "你是一个有趣的SQL助手，可以回答任何问题。",
        }, {
            "role": "user",
            "content": "你好",
        }
    ]

    while True:
        # 发送消息并获取响应
        response_content = await send_message(client, model, messages)

        # 将响应添加到对话历史
        messages.append({"role": "assistant", "content": response_content})

        # 获取用户输入
        user_input = input("\n你: ")
        # 如果用户输入 "exit" 或 "quit"，退出循环
        if user_input.lower() in ["exit", "quit"]:
            break

        # 将用户输入添加到对话历史
        messages.append({"role": "user", "content": user_input})


if __name__ == "__main__":
    # 运行主函数
    asyncio.run(main())
